#include <iostream>
#include <string>
#include <cmath>

#include "MPIIO.h"


#define INDEX(x, y, z)  ((x) + (__int64) ldx * ((y) + (__int64) ldy * (z)))


MPIIO::MPIIO() : buff_size(1000), buff_p(0) {}


void MPIIO::read_yz_block(double* dst, const char* path, const Data3D* data3D, const MPI_TopoCart2D* topo) {
    MPI_File fh;
#ifdef USE_GNU_COMPILER
    char fullpath[255] = {""};
    swt_fullpath(path, fullpath);
    int ret = MPI_File_open(MPI_COMM_WORLD, (char*)fullpath, MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
    if(ret != 0) {
        std::cout << "error with opening file: " << fullpath << "\n";
        exit(1);
    }
#else // USE_GNU_COMPILER
    int ret = MPI_File_open(MPI_COMM_WORLD, path, MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
    if(ret != 0) {
        std::cout << "error with opening file: " << path << "\n";
        exit(1);
    }
#endif // USE_GNU_COMPILER


    int x_size = data3D->x_size, y_size = data3D->y_size;
    int local_blk_size_y = data3D->local_blk_size_y, local_blk_size_z = data3D->local_blk_size_z;
    int BLK_SIZE_Y = data3D->BLK_SIZE_Y, BLK_SIZE_Z = data3D->BLK_SIZE_Z;
    int BLK_LOC_Y = topo->BLK_LOC_Y, BLK_LOC_Z = topo->BLK_LOC_Z;

    __int64 file_disp_bytes = ((__int64)x_size * y_size * (BLK_SIZE_Z * BLK_LOC_Z)
                             + (__int64)BLK_LOC_Y * BLK_SIZE_Y * x_size) * sizeof(double);


    MPI_Datatype file_blk;
    MPI_Type_vector(local_blk_size_z, (__int64)x_size * local_blk_size_y,
                    (__int64)x_size * y_size, MPI_DOUBLE, &file_blk);
    MPI_Type_commit(&file_blk);


    MPI_File_set_view(fh, file_disp_bytes, MPI_DOUBLE, file_blk, "native", MPI_INFO_NULL);

    MPI_File_read(fh, dst, (__int64)x_size * local_blk_size_y * local_blk_size_z,
                  MPI_DOUBLE, MPI_STATUS_IGNORE);

    MPI_File_close(&fh);
}


void MPIIO::blk3D2arr(double* dst, const std::string& out_dir_iden, bool raw, const Data3D* data3D, const MPI_TopoCart2D* topo) {
    std::string whole_path_name = out_dir_iden;

    int halo = data3D->halo;
    int x_size = data3D->x_size, y_size = data3D->y_size;
    int ldx = data3D->ldx, ldy = data3D->ldy;

    int local_blk_size_y = data3D->local_blk_size_y, local_blk_size_z = data3D->local_blk_size_z;
    int BLK_SIZE_Y = data3D->BLK_SIZE_Y, BLK_SIZE_Z = data3D->BLK_SIZE_Z;
    int BLK_LOC_Y = topo->BLK_LOC_Y, BLK_LOC_Z = topo->BLK_LOC_Z;


    MPI_File fh;

    if(raw) {
        whole_path_name += "_raw.dat";
        MPI_File_open(MPI_COMM_WORLD, (char*)whole_path_name.c_str(), MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);

        unsigned __int64 file_disp_bytes = ((__int64)x_size * y_size * (BLK_SIZE_Z * BLK_LOC_Z)
                                            + (__int64)BLK_LOC_Y * BLK_SIZE_Y * x_size) * sizeof(double);

        MPI_Datatype XY_plane;
        MPI_Type_vector(local_blk_size_y, x_size, ldx, MPI_DOUBLE, &XY_plane);
        MPI_Type_commit(&XY_plane);

        MPI_Datatype XYZ_blk;
        MPI_Type_create_hvector(local_blk_size_z, 1, (__int64)ldx * ldy * sizeof(double),
                                XY_plane, &XYZ_blk);
        MPI_Type_commit(&XYZ_blk);


        MPI_Datatype outfile_blk;
        MPI_Type_vector(local_blk_size_z, x_size * local_blk_size_y,
                        x_size * y_size, MPI_DOUBLE, &outfile_blk);
        MPI_Type_commit(&outfile_blk);


        MPI_File_set_view(fh, file_disp_bytes, XYZ_blk, outfile_blk, "native", MPI_INFO_NULL);
        MPI_File_write(fh, dst + (__int64)ldx * ldy * halo + ldx * halo, 1,
                       XYZ_blk, MPI_STATUS_IGNORE);

    } else {
        //! this operation will change dst[i] to abs(dst[i])
        whole_path_name += ".dat";
        MPI_File_open(MPI_COMM_WORLD, (char*)whole_path_name.c_str(), MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);

        // process wave-field to 0-255 values
        __int64 blk_domain_size = data3D->blk_domain_size;

        for(__int64 i = 0; i < blk_domain_size; i++)
            dst[i] = fabs(dst[i]);
        double* max_arr = new double[topo->comm_siz];
        double temp_max = 0;

        for(__int64 i = 0; i < blk_domain_size; i++)
            if(dst[i] > temp_max)
                temp_max = dst[i];
        max_arr[topo->proc_id] = temp_max;

        std::cout << "proc " << topo->proc_id << " : local_max = " << temp_max << "\n";

        MPI_Barrier(MPI_COMM_WORLD);

        MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, max_arr, 1, MPI_DOUBLE, MPI_COMM_WORLD);


        for(int ii = 0; ii < topo->comm_siz; ii++)       // use temp_max to store the max value of dst across all processes
            if(max_arr[ii] > temp_max)
                temp_max = max_arr[ii];

        // std::cout << "proc " << proc_id << " : global_max = " << temp_max << ".\n";
        short* values = new short[blk_domain_size];
        for(__int64 i = 0; i < blk_domain_size; i++)
            values[i] = short(dst[i] / temp_max * 255);


        unsigned __int64 file_disp_bytes = ((__int64)x_size * y_size * (BLK_SIZE_Z * BLK_LOC_Z)
                                          + (__int64)BLK_LOC_Y * BLK_SIZE_Y * x_size) * sizeof(short);

        MPI_Datatype XY_plane;
        MPI_Type_vector(local_blk_size_y, x_size, ldx, MPI_SHORT, &XY_plane);
        MPI_Type_commit(&XY_plane);

        MPI_Datatype XYZ_blk;
        MPI_Type_create_hvector(local_blk_size_z, 1, (__int64)ldx * ldy * sizeof(short),
                                XY_plane, &XYZ_blk);
        MPI_Type_commit(&XYZ_blk);


        MPI_Datatype outfile_blk;
        MPI_Type_vector(local_blk_size_z, (__int64)x_size * local_blk_size_y,
                        (__int64)x_size * y_size, MPI_SHORT, &outfile_blk);
        MPI_Type_commit(&outfile_blk);


        MPI_File_set_view(fh, file_disp_bytes, XYZ_blk, outfile_blk, "native", MPI_INFO_NULL);
        MPI_File_write(fh, values + (__int64)ldx * ldy * halo + ldx * halo, 1,
                       XYZ_blk, MPI_STATUS_IGNORE);

        delete[] values;
        delete[] max_arr;
    }

    MPI_File_close(&fh);
}


void MPIIO::blk3Dslice_x(double* dst, const std::string& name, bool raw, const Data3D* data3D, const MPI_TopoCart2D* topo, int target_x) {
//! this function will not change the values of *dst

    int halo = data3D->halo;
    int y_size = data3D->y_size;
    int ldx = data3D->ldx, ldy = data3D->ldy;

    int local_blk_size_y = data3D->local_blk_size_y, local_blk_size_z = data3D->local_blk_size_z;
    int BLK_SIZE_Y = data3D->BLK_SIZE_Y, BLK_SIZE_Z = data3D->BLK_SIZE_Z;
    int BLK_LOC_Y = topo->BLK_LOC_Y, BLK_LOC_Z = topo->BLK_LOC_Z;


    int dst_x = target_x + halo;


    MPI_File fh;

    MPI_File_open(MPI_COMM_WORLD, (char *)name.c_str(), MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);

    if(raw) {
        double *buffer = new double[(__int64) local_blk_size_y * local_blk_size_z];
        for(int k = 0; k < local_blk_size_z; k++) {
            __int64 disp = (__int64) k * local_blk_size_y;
            int dst_k = k + halo;

            for(int j = 0; j < local_blk_size_y; j++)
                buffer[disp + j] = dst[INDEX(dst_x, j + halo, dst_k)];
        }

        MPI_Datatype file_blk;
        MPI_Type_vector(local_blk_size_z, local_blk_size_y, y_size, MPI_DOUBLE, &file_blk);
        MPI_Type_commit(&file_blk);

        MPI_File_set_view(fh, ((__int64) BLK_LOC_Z * y_size * BLK_SIZE_Z + BLK_SIZE_Y * BLK_LOC_Y) * sizeof(double),
                          MPI_DOUBLE, file_blk, "native", MPI_INFO_NULL);

        MPI_File_write(fh, buffer, (__int64) local_blk_size_y * local_blk_size_z, MPI_DOUBLE, MPI_STATUS_IGNORE);

        delete[] buffer;
    } else {
        // collect a slice and process it to 0-255 values

        double *max_arr = new double[topo->comm_siz];
        double *buffer = new double[(__int64) local_blk_size_y * local_blk_size_z];

        double temp_max = 0;

        for(int k = 0; k < local_blk_size_z; k++) {
            __int64 disp = (__int64) k * local_blk_size_y;
            int dst_k = k + halo;
            for(int j = 0; j < local_blk_size_y; j++) {
                buffer[disp + j] = fabs(dst[INDEX(dst_x, j + halo, dst_k)]);
                if(buffer[disp + j] > temp_max)
                    temp_max = buffer[disp + j];
            }
        }

        max_arr[topo->proc_id] = temp_max;



        MPI_Barrier(MPI_COMM_WORLD);

        MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, max_arr, 1, MPI_DOUBLE, MPI_COMM_WORLD);


        for(int i = 0; i < topo->comm_siz; i++)       // use temp_max to store the max value of dst across all processes
            if(max_arr[i] > temp_max)
                temp_max = max_arr[i];



        short *slice = new short[(__int64) local_blk_size_y * local_blk_size_z];
        for(__int64 k = 0; k < local_blk_size_z; k++) {
            __int64 disp = (__int64) k * local_blk_size_y;
            for(int j = 0; j < local_blk_size_y; j++)
                slice[disp + j] = short(buffer[disp + j] / temp_max * 255);
        }


        MPI_Datatype short_file_blk;
        MPI_Type_vector(local_blk_size_z, local_blk_size_y, y_size, MPI_SHORT, &short_file_blk);
        MPI_Type_commit(&short_file_blk);

        MPI_File_set_view(fh, ((__int64) BLK_LOC_Z * y_size * BLK_SIZE_Z + BLK_SIZE_Y * BLK_LOC_Y) * sizeof(short),
                          MPI_SHORT, short_file_blk, "native", MPI_INFO_NULL);

        MPI_File_write(fh, slice, (__int64) local_blk_size_y * local_blk_size_z, MPI_SHORT, MPI_STATUS_IGNORE);


        delete[] slice;
        delete[] buffer;
        delete[] max_arr;
    }

    MPI_File_close(&fh);

}


void MPIIO::blk3Dslice_y(double* dst, const std::string& name, bool raw, const Data3D* data3D, const MPI_TopoCart2D* topo, int target_y) {
    //! this function will not change the values of *dst

    int halo = data3D->halo;
    int x_size = data3D->x_size, y_size = data3D->y_size;
    int ldx = data3D->ldx, ldy = data3D->ldy;

    int local_blk_size_y = data3D->local_blk_size_y, local_blk_size_z = data3D->local_blk_size_z;
    int BLK_SIZE_Y = data3D->BLK_SIZE_Y, BLK_SIZE_Z = data3D->BLK_SIZE_Z;
    int BLK_LOC_Y = topo->BLK_LOC_Y, BLK_LOC_Z = topo->BLK_LOC_Z;

    int out_blk_y_loc = target_y / BLK_SIZE_Y;


    MPI_Group group_world, group_out;
    MPI_Comm comm_loc_y;
    MPI_Comm_group(MPI_COMM_WORLD, &group_world);

    int* members = new int[topo->NUM_BLK_Z];
    for(int i = 0; i < topo->NUM_BLK_Z; i++)
        members[i] = i * topo->NUM_BLK_Y + BLK_LOC_Y;

    MPI_Group_incl(group_world, topo->NUM_BLK_Z, members, &group_out);
    MPI_Comm_create(MPI_COMM_WORLD, group_out, &comm_loc_y);
    delete[] members;

    if(out_blk_y_loc == BLK_LOC_Y) {

        int dst_y = target_y - out_blk_y_loc * BLK_SIZE_Y + halo;

        MPI_File fh;

        MPI_File_open(comm_loc_y, (char*)name.c_str(), MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);

        if(raw) {
            double* buffer = new double[(__int64) x_size * local_blk_size_z];
            for(int k = 0; k < local_blk_size_z; k++) {
                __int64 disp = (__int64) k * x_size;
                int dst_k = k + halo;

                for(int i = 0; i < x_size; i++)
                    buffer[disp + i] = dst[INDEX(i + halo, dst_y, dst_k)];
            }

            MPI_File_set_view(fh, (__int64) BLK_LOC_Z * x_size * BLK_SIZE_Z * sizeof(double),
                              MPI_DOUBLE, MPI_DOUBLE, "native", MPI_INFO_NULL);

            MPI_File_write(fh, buffer, (__int64) x_size * local_blk_size_z, MPI_DOUBLE, MPI_STATUS_IGNORE);

            delete[] buffer;
        } else {
            // collect a slice and process it to 0-255 values
            int new_comm_siz;
            int new_proc_id;
            MPI_Comm_size(comm_loc_y, &new_comm_siz);
            MPI_Comm_rank(comm_loc_y, &new_proc_id);
            double* max_arr = new double[new_comm_siz];
            double* buffer = new double[(__int64) x_size * local_blk_size_z];

            double temp_max = 0;
            for(int k = 0; k < local_blk_size_z; k++) {
                __int64 disp = (__int64) k * x_size;

                int dst_k = k + halo;
                for(int i = 0; i < x_size; i++) {
                    buffer[disp + i] = fabs(dst[INDEX(i + halo, dst_y, dst_k)]);
                    if(buffer[disp + i] > temp_max)
                        temp_max = buffer[disp + i];
                }
            }

            max_arr[new_proc_id] = temp_max;


            MPI_Barrier(comm_loc_y);

            MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, max_arr, 1, MPI_DOUBLE, comm_loc_y);


            for(int i = 0; i < new_comm_siz; i++)       // use temp_max to store the max value of dst across all processes
                if(max_arr[i] > temp_max)
                    temp_max = max_arr[i];


            short* slice = new short[(__int64) x_size * local_blk_size_z];
            for(int k = 0; k < local_blk_size_z; k++) {
                __int64 disp = (__int64) k * x_size;
                for(int i = 0; i < x_size; i++)
                    slice[disp + i] = short(buffer[disp + i] / temp_max * 255);
            }


            MPI_File_set_view(fh,(__int64) BLK_LOC_Z * x_size * BLK_SIZE_Z * sizeof(short),
                              MPI_SHORT, MPI_SHORT, "native", MPI_INFO_NULL);

            MPI_File_write(fh, slice, (__int64) x_size * local_blk_size_z, MPI_SHORT, MPI_STATUS_IGNORE);


            delete[] slice;
            delete[] buffer;
            delete[] max_arr;
        }

        MPI_File_close(&fh);
    }

    MPI_Comm_free(&comm_loc_y);
}


void MPIIO::blk3Dslice_z(double* dst, const std::string& name, bool raw, const Data3D* data3D, const MPI_TopoCart2D* topo, int target_z) {
    //! this function will not change the values of *dst

    int halo = data3D->halo;
    int x_size = data3D->x_size, y_size = data3D->y_size;
    int ldx = data3D->ldx, ldy = data3D->ldy;

    int local_blk_size_y = data3D->local_blk_size_y, local_blk_size_z = data3D->local_blk_size_z;
    int BLK_SIZE_Y = data3D->BLK_SIZE_Y, BLK_SIZE_Z = data3D->BLK_SIZE_Z;
    int BLK_LOC_Y = topo->BLK_LOC_Y, BLK_LOC_Z = topo->BLK_LOC_Z;

    int out_blk_z_loc = target_z / BLK_SIZE_Z;


    MPI_Group group_world, group_out;
    MPI_Comm comm_loc_z;
    MPI_Comm_group(MPI_COMM_WORLD, &group_world);

    int* members = new int[topo->NUM_BLK_Y];
    for(int i = 0; i < topo->NUM_BLK_Y; i++)
        members[i] = i + topo->NUM_BLK_Y * BLK_LOC_Z;

    MPI_Group_incl(group_world, topo->NUM_BLK_Y, members, &group_out);
    MPI_Comm_create(MPI_COMM_WORLD, group_out, &comm_loc_z);
    delete[] members;

    if(out_blk_z_loc == BLK_LOC_Z) {

        int dst_z = target_z - out_blk_z_loc * BLK_SIZE_Z + halo;


        MPI_File fh;

        MPI_File_open(comm_loc_z, (char*)name.c_str(), MPI_MODE_CREATE | MPI_MODE_WRONLY, MPI_INFO_NULL, &fh);


        if(raw) {
            double *buffer = new double[(__int64) x_size * local_blk_size_y];
            for(int j = 0; j < local_blk_size_y; j++) {
                __int64 disp = (__int64) j * x_size;

                int dst_j = j + halo;
                for (int i = 0; i < x_size; i++)
                    buffer[disp + i] = dst[INDEX(i + halo, dst_j, dst_z)];
            }

            MPI_File_set_view(fh, (__int64) BLK_LOC_Y * x_size * BLK_SIZE_Y * sizeof(double),
                              MPI_DOUBLE, MPI_DOUBLE, "native", MPI_INFO_NULL);

            MPI_File_write(fh, buffer, (__int64) x_size * local_blk_size_y, MPI_DOUBLE, MPI_STATUS_IGNORE);

            delete[] buffer;
        } else {
            // collect a slice and process it to 0-255 values
            int new_comm_siz;
            int new_proc_id;
            MPI_Comm_size(comm_loc_z, &new_comm_siz);
            MPI_Comm_rank(comm_loc_z, &new_proc_id);
            double *max_arr = new double[new_comm_siz];
            double *buffer = new double[(__int64) x_size * local_blk_size_y];

            double temp_max = 0;
            for(int j = 0; j < local_blk_size_y; j++) {
                __int64 disp = (__int64) j * x_size;
                int dst_j = j + halo;
                for(int i = 0; i < x_size; i++) {
                    buffer[disp + i] = fabs(dst[INDEX(i + halo, dst_j, dst_z)]);
                    if (buffer[disp + i] > temp_max)
                        temp_max = buffer[disp + i];
                }
            }

            max_arr[new_proc_id] = temp_max;


            MPI_Barrier(comm_loc_z);

            MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, max_arr, 1, MPI_DOUBLE, comm_loc_z);


            for(int i = 0; i < new_comm_siz; i++)       // use temp_max to store the max value of dst across all processes
                if(max_arr[i] > temp_max)
                    temp_max = max_arr[i];


            short *slice = new short[(__int64) x_size * local_blk_size_y];
            for(int j = 0; j < local_blk_size_y; j++) {
                __int64 disp = (__int64) j * x_size;
                for(int i = 0; i < x_size; i++)
                    slice[disp + i] = short(buffer[disp + i] / temp_max * 255);
            }


            MPI_File_set_view(fh,(__int64) BLK_LOC_Y * x_size * BLK_SIZE_Y * sizeof(short),
                              MPI_SHORT, MPI_SHORT, "native", MPI_INFO_NULL);

            MPI_File_write(fh, slice, (__int64) x_size * local_blk_size_y, MPI_SHORT, MPI_STATUS_IGNORE);


            delete[] slice;
            delete[] buffer;
            delete[] max_arr;
        }

        MPI_File_close(&fh);
    }

    MPI_Comm_free(&comm_loc_z);
}


void MPIIO::out_coord_val(const double* dst, double* buffer, const std::vector<OutputPtInfo>& out_pt_info) {
    for(int i = 0; i < out_pt_info.size(); i++)
        buffer[i * buff_size + buff_p] = dst[out_pt_info[i].pos];

    buff_p++;
    if(buff_p == buff_size)
        flush_buffer(buffer, out_pt_info);
}


void MPIIO::flush_buffer(double* buffer, const std::vector<OutputPtInfo>& out_pt_info) {
    for(int i = 0; i < out_pt_info.size(); i++) {
        // handle the file with normal fopen
        FILE *fp = fopen(out_pt_info[i].name.c_str(), "ab");
        if(!fp) {
            std::cout << "error with opening file: " << out_pt_info[i].name << "\n";
            exit(1);
        }

        fwrite(&buffer[i * buff_size], sizeof(double), buff_p, fp);
        fclose(fp);
    }
    buff_p = 0;
}
